// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/dwrf/writer/statistics_builder.h>

namespace kumo::pollux::dwrf {

class PhysicalSizeAggregator {
 public:
  explicit PhysicalSizeAggregator(PhysicalSizeAggregator* parent = nullptr)
      : parent_{parent} {}

  virtual ~PhysicalSizeAggregator() = default;

  virtual void recordSize(const DwrfStreamIdentifier& id, uint64_t streamSize) {
    result_ += streamSize;
    if (parent_ != nullptr) {
      parent_->recordSize(id, streamSize);
    }
  }

  uint64_t getResult() {
    return result_;
  }

 private:
  PhysicalSizeAggregator* const parent_;
  uint64_t result_{0};
};

class MapPhysicalSizeAggregator : public PhysicalSizeAggregator {
 public:
  explicit MapPhysicalSizeAggregator(PhysicalSizeAggregator* parent = nullptr)
      : PhysicalSizeAggregator{parent} {}

  virtual ~MapPhysicalSizeAggregator() = default;

  void recordSize(const DwrfStreamIdentifier& id, uint64_t streamSize)
      override {
    PhysicalSizeAggregator::recordSize(id, streamSize);
    if (!mapStatsBuilder_) {
      return;
    }
    if (id.encodingKey().sequence() != 0) {
      mapStatsBuilder_->incrementSize(
          sequenceToKey_.at(id.encodingKey().sequence()), streamSize);
    }
  }

  void prepare(
      melon::F14FastMap<uint32_t, const proto::KeyInfo&> sequenceToKey,
      MapStatisticsBuilder* mapStatsBuilder) {
    sequenceToKey_ = std::move(sequenceToKey);
    mapStatsBuilder_ = mapStatsBuilder;
  }

 private:
  melon::F14FastMap<uint32_t, const proto::KeyInfo&> sequenceToKey_;
  MapStatisticsBuilder* mapStatsBuilder_{nullptr};
};

} // namespace kumo::pollux::dwrf
